# 从 transformers 库中导入自动模型类和分词器类，用于加载大语言模型及其对应的 tokenizer
from transformers import AutoModelForCausalLM, AutoTokenizer

# 导入 PyTorch 深度学习框架
import torch

# 设置随机种子为 0，确保实验的可复现性（即每次运行结果一致）
torch.manual_seed(0)

# 指定本地模型路径，MiniCPM-4.1-8B 模型存储在该目录下
path = '/root/autodl-tmp/llm_models/MiniCPM4.1-8B'

# 指定模型运行设备为 GPU（CUDA）
device = "cuda"

# 使用 AutoTokenizer 加载指定路径下的分词器（tokenizer），用于将文本转换为模型可处理的 token ID
tokenizer = AutoTokenizer.from_pretrained(path)

# 加载预训练的因果语言模型（AutoModelForCausalLM）
# 参数说明：
# - torch_dtype=torch.bfloat16: 使用 bfloat16 精度以节省显存并提升推理速度
# - device_map=device: 将模型加载到指定设备（这里是 GPU）
# - trust_remote_code=True: 允许执行模型仓库中的自定义代码（某些模型需要）
model = AutoModelForCausalLM.from_pretrained(
    path, 
    torch_dtype=torch.bfloat16, 
    device_map=device, 
    trust_remote_code=True
)

# 定义对话历史消息列表，模拟用户提问
# role 为 "user" 表示这是用户输入的内容
messages = [
    {"role": "user", "content": "写一篇人工智能为题的文章。50字内。"},
]

# 使用 tokenizer 的 apply_chat_template 方法将对话消息格式化为模型所需的输入文本
# 参数说明：
# - tokenize=False: 不进行分词，只返回原始文本
# - add_generation_prompt=True: 在末尾添加生成提示符（如 "Assistant:"），告诉模型开始生成回复
# - enable_thinking=False: MiniCPM4.1支持混合推理模式，可以在深度推理模式和非推理模式下使用。
prompt_text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False   # 打开关闭推理模式
)

# 将格式化后的提示文本转换为模型输入张量（Tensor）
# return_tensors="pt" 表示返回 PyTorch 张量
# .to(device) 将输入数据移动到 GPU 上
model_inputs = tokenizer([prompt_text], return_tensors="pt").to(device)

# 调用模型的 generate 方法进行文本生成
# 参数说明：
# - max_new_tokens=32768: 最多生成 32768 个新 token（注意：这个值非常大，可能造成显存溢出或耗时过长）
# - top_p=0.95: 使用核采样（nucleus sampling），保留累计概率前 95% 的词汇
# - temperature=0.6: 控制生成文本的随机性，值越低越确定，越高越随机
model_outputs = model.generate(
    **model_inputs,
    max_new_tokens=32768,
    top_p=0.95,
    temperature=0.6
)

# 提取生成的 token ID
# 原始输出包含输入 prompt 的 token，因此需要切片去除，只保留新生成的部分
# len(model_inputs[i]) 这里写法有误，应为 len(model_inputs['input_ids'][i])
# 正确做法是：对每个样本，从输出中去掉输入部分，保留新生成的 token
# 注意：此处代码存在潜在 bug，应改为如下形式：
# output_token_ids = [model_outputs[i][len(model_inputs['input_ids'][i]):] for i in range(len(model_inputs['input_ids']))]
output_token_ids = [
    model_outputs[i][len(model_inputs[i]):] for i in range(len(model_inputs['input_ids']))
]

# 将生成的 token ID 解码回可读文本
# skip_special_tokens=True: 忽略特殊标记（如 <s>, </s>, <pad> 等）
# [0] 取出第一个（也是唯一一个）生成结果
responses = tokenizer.batch_decode(output_token_ids, skip_special_tokens=True)[0]

# 打印最终生成的回复文本
print(responses)